import pygame
import numpy as np
import math
from random import randint, random, shuffle

FPS = 60
WINDOW_SIZE = (800, 600)

class Ball:
  def __init__(self, position, mass, velocity=None) -> None:
    self.position = position
    self.set_mass(mass)
    if velocity is not None:
      self.velocity = velocity
    else:
      coeff = 100000 / mass
      angle = random() * (2 * math.pi)
      self.velocity = np.array([coeff * math.cos(angle), coeff * math.sin(angle)])
    color = [255, 7, randint(0, 255)]
    shuffle(color)
    self.color = tuple(color)

  def set_mass(self, mass):
    self.mass = mass
    self.radius = math.sqrt(mass / math.pi)

  def update(self, balls, delta_time):

    # 边界碰撞反弹 + 防止卡在边界
    if self.position[0] - self.radius < 0:
      self.velocity[0] = -self.velocity[0]
      self.position[0] = self.radius
    if self.position[0] + self.radius > WINDOW_SIZE[0]:
      self.velocity[0] = -self.velocity[0]
      self.position[0] = WINDOW_SIZE[0] - self.radius
    if self.position[1] - self.radius < 0:
      self.velocity[1] = -self.velocity[1]
      self.position[1] = self.radius
    if self.position[1] + self.radius > WINDOW_SIZE[1]:
      self.velocity[1] = -self.velocity[1]
      self.position[1] = WINDOW_SIZE[1] - self.radius

    for ball in balls:
      if self is ball:
        continue

      # 碰撞检测
      distance = math.hypot(ball.position[0] - self.position[0], ball.position[1] - self.position[1])
      total_radius = self.radius + ball.radius
      if distance >= total_radius:
        continue

      # 防止两个球卡在一起
      angle1 = math.atan2(self.position[1] - ball.position[1], self.position[0] - ball.position[0])
      self.position[0] += math.cos(angle1) * (total_radius - distance)
      self.position[1] += math.sin(angle1) * (total_radius - distance)

      # 计算完全弹性碰撞速度
      total_mass = self.mass + ball.mass
      new_self_velocity = (((self.mass - ball.mass) * self.velocity) + ((2 * ball.mass) * ball.velocity)) / total_mass
      new_ball_velocity = (((ball.mass - self.mass) * ball.velocity) + ((2 * self.mass) * self.velocity)) / total_mass

      self.velocity = new_self_velocity
      ball.velocity = new_ball_velocity

      # 计算非弹性碰撞速度
      # final_velocity = (ball.mass * ball.velocity - self.mass * self.velocity) / total_mass

      # 计算完全非弹性碰撞速度
      # final_velocity = (self.mass * self.velocity + ball.mass * ball.velocity) / total_mass

      # self.velocity = new_self_velocity
      # ball.velocity = new_ball_velocity

    self.position += self.velocity * delta_time

  def draw(self, surface):
    pygame.draw.circle(surface, self.color, (self.position[0], self.position[1]), self.radius)


balls = [Ball(np.random.random((2)) * np.array(WINDOW_SIZE), randint(100, 2000)) for _ in range(50)]

pygame.init()
scene = pygame.display.set_mode(WINDOW_SIZE)
clock = pygame.time.Clock()
font = pygame.font.SysFont('方正粗黑宋简体', 24)
running = True
while running:
  clock.tick(FPS)

  for event in pygame.event.get():
    if event.type == pygame.QUIT:
      running = False

  delta_time = clock.get_time() / 1000.0

  for ball in balls:
    ball.update(balls, delta_time)

  scene.fill((255, 255, 255))

  for ball in balls:
    ball.draw(scene)

  text = font.render(f'fps {math.floor(clock.get_fps())}', True, 'green')
  rect = text.get_rect(topleft=(10, 10))
  scene.blit(text, rect)

  pygame.display.flip()

pygame.quit()
